#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 19 10:50:14 2022

@author: liangy
E-mail: liangy19@mails.tsinghua.edu.cn

For galaxy radial profiles, both spherical and cylinderical

Unit System
The unit system applied here is:
     distance         --> kiloparsec 
     angular distance --> arcsec
     mass             --> solar mass
     time             --> Gyr
     velocity         --> km/s
"""

import math
import numpy as np
from scipy import special
from scipy.integrate import quad,dblquad

# for two power law projection
def aux_integral_1(s,s0,a,b):
    return s/(math.sqrt(s*s-s0*s0)*math.pow(s,a)*math.pow(1+s,b-a))

def aux_integral_1d(s,s0,a,b):
    return (a+b*s)/(math.sqrt(s*s-s0*s0)*math.pow(s,a+1)*math.pow(1+s,b-a+1))

def aux_integral_2(s,s0,a,b):
    return s0*s/(math.sqrt(s*s-s0*s0)*math.pow(s,a)*math.pow(1+s,b-a))

# for sersic profile deprojection
def aux_integral_3(s,s0,b,m):
    return math.pow(s,1/m-1)*math.exp(-b*math.pow(s,1/m))/math.sqrt(s*s-s0*s0)

def aux_integral_4(s,s0,b,m):
    return s0*s0*math.pow(s,1/m-1)*math.exp(-b*math.pow(s,1/m))/math.sqrt(s*s-s0*s0)

# for einasto profile projection
def aux_integral_5(s,s0,d,n):
    return s*math.exp(-d*math.pow(s,1/n))/math.sqrt(s*s-s0*s0)

def aux_integral_6(s,s0,d,n):
    return s0*s*math.exp(-d*math.pow(s,1/n))/math.sqrt(s*s-s0*s0)

# for analytic hernquist projection
def X(s):
    if type(s) == np.ndarray:
        Xs = np.zeros_like(s)
        Xs[s <= 1.0] = np.arccosh(1/s[s<=1.0])/np.sqrt(1.-s[s<=1.0]*s[s<=1.0])
        Xs[s > 1.0] = np.arccos(1/s[s>1.0])/np.sqrt(s[s>1.0]*s[s>1.0]-1.)
        return Xs
    else:
        if s <= 1.0:
            return np.arccosh(1/s)/np.sqrt(1.-s*s)
        if s > 1.0:
            return np.arccos(1/s)/np.sqrt(s*s-1.)

class sph_profile():
    """
    To define the radial profile function in spherical or cylinderical coordinate
    """
    def __init__(self,radial_coordinate,dim = '3D'):
        # radial_coordinate: float or np.ndarray, unit --> kpc
        self.r   = radial_coordinate
        self.lgr = np.log10(radial_coordinate)
        self.dim = dim    # 3D density or 2D surface density 
        
    def power_law(self,rho_0,r_s,a):
        # power law 3D density profile
        if self.dim == '3D':
            s = self.r/r_s    # scaled by r_s(scale radius) 
            return rho_0/np.power(s,a) 
        # power law 2D projected surface density profile
        if self.dim == '2D':
            s = self.r/r_s    # scaled by r_s(scale radius) 
            return rho_0*r_s*special.beta((a-1)/2,0.5)*np.power(s,1.0-a)
        
    def d_power_law(self,rho_0,r_s,a):
        # radial gradient of power law 3D density profile
        if self.dim == '3D':
            s = self.r/r_s    # scaled by r_s(scale radius) 
            return -a*rho_0/r_s/np.power(s,a+1)
        # radial gradient of power law 2D projected surface density profile
        if self.dim == '2D':
            s = self.r/r_s    # scaled by r_s(scale radius) 
            return (1.0-a)*rho_0*special.beta((a-1)/2,0.5)/np.power(s,a)
        
    def M_power_law(self,rho_0,r_s,a):
        # 3D enclosed mass of power law density profile
        if self.dim == '3D':
            s = self.r/r_s    # scaled by r_s(scale radius) 
            return (4*np.pi*rho_0*r_s*r_s*r_s)*np.power(s,3.0-a)/(3.0-a)
        # 2D enclosed mass of power law projected surface density profile
        # 1< a < 3
        if self.dim == '2D':
            s = self.r/r_s    # scaled by r_s(scale radius) 
            return 2*np.pi*rho_0*r_s*r_s*r_s*special.beta((a-1)/2,0.5)*np.power(s,3.0-a)/(3.0-a)
        
    def dM_power_law(self,rho_0,r_s,a):
        # radial gradient of 3D enclosed mass of power law density profile
        if self.dim == '3D':
            s = self.r/r_s    # scaled by r_s(scale radius
            return (4*np.pi*rho_0*r_s*r_s)*np.power(s,2.0-a)
        if self.dim == '2D':
            s = self.r/r_s
            return 2*np.pi*rho_0*r_s*r_s*special.beta((a-1)/2,0.5)*np.power(s,2.0-a)
        
    def two_power_law(self,rho_0,r_s,a,b):
        # two power law 3D density profile
        if self.dim == '3D':
            s = self.r/r_s
            return rho_0/np.power(s,a)/np.power(1+s,b-a)
        # two power law 2D projected surface density profile 
        if self.dim == '2D':
            s = self.r/r_s
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = quad(aux_integral_1,s[i],np.inf,args = (s[i],a,b))
            else:
                I,Ierr = quad(aux_integral_1,s,np.inf,args = (s,a,b))
            return 2*r_s*rho_0*I
        
    def d_two_power_law(self,rho_0,r_s,a,b):
        # radial gradient of two power law 3D density profile
        if self.dim == '3D':
            s = self.r/r_s
            return -rho_0/r_s/np.power(s,a+1)/np.power(1+s,b-a+1)*(a+b*s)
        # radial gradient of two power law 2D projected surface density profile
        if self.dim == '2D':
            s = self.r/r_s
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = quad(aux_integral_1d,s[i],np.inf,args = (s[i],a,b))
            else:
                I,Ierr = quad(aux_integral_1d,s,np.inf,args = (s,a,b))
            return -2*rho_0*s*I
        
    def M_two_power_law(self,rho_0,r_s,a,b):
        # 3D enclosed mass of two power law density profile
        # a < 3, b > 3
        if self.dim == '3D':
            if b == 3:
                b += 3e-7    # perturbation to avoid singularities in gamma and betainc
            s = self.r/r_s
            return 4*np.pi*rho_0*r_s*r_s*r_s*special.beta(3-a,b-3)*special.betainc(3-a,b-3,s/(1+s))
        # 2D enclosed mass of two power law surface density profile
        if self.dim == '2D':
            def aux_lim(s0):
                return s0
            s = self.r/r_s
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = dblquad(aux_integral_2,0,s[i],aux_lim,np.inf,args = (a,b))
            else:
                I,Ierr = dblquad(aux_integral_2,0,s,aux_lim,np.inf,args = (a,b))
            return 4*np.pi*r_s*r_s*r_s*rho_0*I
        
    def dM_two_power_law(self,rho_0,r_s,a,b):
        # radial gradient of 3D enclosed mass of two power law density profile
        if self.dim == '3D':
            s = self.r/r_s
            return 4*np.pi*r_s*r_s*rho_0/np.power(s,a-2)/np.power(1+s,b-a)
        # radial gradient of 2D enclosed mass of two power law surface density profile
        if self.dim == '2D':
            s = self.r/r_s
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = quad(aux_integral_1,s[i],np.inf,args = (s[i],a,b))
            else:
                I,Ierr = quad(aux_integral_1,s,np.inf,args = (s,a,b))
            return 4*np.pi*s*r_s*r_s*rho_0*I
        
    def sersic(self,L0,Re,m):
        # sersic 2D surface brightness profile
        if self.dim == '2D':
            s = self.r/Re
            b = special.gammaincinv(2*m,0.5)
            return np.power(b,2*m)*L0/(2*np.pi*m*special.gamma(2*m)*Re*Re)*np.exp(-b*np.power(s,1/m))
        # deprojected sersic 3D luminosity profile
        if self.dim == '3D':
            s = self.r/Re
            b = special.gammaincinv(2*m,0.5)
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = quad(aux_integral_3,s[i],np.inf,args = (s[i],b,m))
            else:
                I,Ierr = quad(aux_integral_3,s,np.inf,args = (s,b,m))
            return np.power(b,2*m+1)*L0/(2*np.pi*np.pi*m*m*special.gamma(2*m)*Re*Re*Re)*I
        
    def d_sersic(self,L0,Re,m):
        # radial gradient of sersic 2D surface brightness profile
        if self.dim == '2D':
            s = self.r/Re
            b = special.gammaincinv(2*m,0.5)
        return -np.power(b,2*m+1)*L0/(2*np.pi*m*m*special.gamma(2*m)*Re*Re*Re)*np.exp(-b*np.power(s,1/m))*np.power(s,1/m-1)
    
    def M_sersic(self,L0,Re,m):
        # 2D enclosed light of sersic surface brightness profile
        if self.dim == '2D':
            s = self.r/Re
            b = special.gammaincinv(2*m,0.5)
            return L0*special.gammainc(2*m,b*np.power(s,1/m))
        # 3D enclosed light of deprojected sersic luminosity profile
        if self.dim == '3D':
            def aux_lim(s0):
                return s0
            s = self.r/Re
            b = special.gammaincinv(2*m,0.5)
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = dblquad(aux_integral_4,0,s[i],aux_lim,np.inf,args = (b,m))
            else:
                I,Ierr = dblquad(aux_integral_4,0,s,aux_lim,np.inf,args = (b,m))
            return 2*np.power(b,2*m+1)*L0/(np.pi*m*m*special.gamma(2*m))*I
        
    def dM_sersic(self,L0,Re,m):
        # radial gradient of sersic 2D surface brightness profile
        if self.dim == '2D':
            s = self.r/Re
            b = special.gammaincinv(2*m,0.5)
            return np.power(b,2*m)*L0/(m*special.gamma(2*m)*Re)*s*np.exp(-b*np.power(s,1/m))
        
    def einasto(self,M0,r_h,n):
        # einasto 3D dark matter density profile
        if self.dim == '3D':
            s = self.r/r_h
            d = special.gammaincinv(3*n,0.5)
            return np.power(d,3*n)*M0/(4*np.pi*n*special.gamma(3*n)*r_h*r_h*r_h)*np.exp(-d*np.power(s,1/n))
        # einasto 2D surface dark matter density profile
        if self.dim == '2D':
            s = self.r/r_h
            d = special.gammaincinv(3*n,0.5)
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = quad(aux_integral_5,s[i],np.inf,args = (s[i],d,n))
            else:
                I,Ierr = quad(aux_integral_5,s,np.inf,args = (s,d,n))
            return np.power(d,3*n)*M0/(2*np.pi*n*special.gamma(3*n)*r_h*r_h)*I
        
        
    def M_einasto(self,M0,r_h,n):
        # 3D enclosed mass of einasto density profile
        if self.dim == '3D':
            s = self.r/r_h
            d = special.gammaincinv(3*n,0.5)
            return M0*special.gammainc(3*n,d*np.power(s,1/n))
        if self.dim == '2D':
            def aux_lim(s0):
                return s0
            s = self.r/r_h
            d = special.gammaincinv(3*n,0.5)
            if type(s) == np.ndarray:
                I    = np.zeros_like(s)
                Ierr = np.zeros_like(s)
                for i in range(len(s)):
                    I[i],Ierr[i] = dblquad(aux_integral_6,0,s[i],aux_lim,np.inf,args = (d,n))
            else:
                I,Ierr = dblquad(aux_integral_6,0,s,aux_lim,np.inf,args = (d,n))
            return np.power(d,3*n)*M0/(n*special.gamma(3*n))*I
    
    def gaussian(self,M0,sigma):
        # gaussian 3D density profile
        if self.dim == '3D':
            s = self.r/sigma
            return M0/(np.sqrt(2*np.pi)*sigma)**3*np.exp(-0.5*s*s)
        # gaussian 2D surface brightness profile
        if self.dim == '2D':
            s = self.r/sigma
            return M0/(np.sqrt(2*np.pi)*sigma)**2*np.exp(-0.5*s*s)
        
    def M_gaussian(self,M0,sigma):
        # 3D enclosed mass of gaussian profile
        if self.dim == '3D':
            s = self.r/sigma
            return M0*(special.erf(s/np.sqrt(2))-2*s/np.sqrt(2*np.pi)*np.exp(-0.5*s*s))
        # 2D enclosed mass of gaussian surface brightness profile
        if self.dim == '2D':
            s = self.r/sigma
            return M0*(1-np.exp(-0.5*s*s))
    
    # the specific routine for hernquist projection
    def hernquist_2D(self,rho_0,r_s):
        s = self.r/r_s
        if type(s) == np.ndarray:
            Sigma = np.zeros_like(s)
            singularity = np.abs(s - 1.0) <= 1e-4
            ordinary    = np.abs(s - 1.0) > 1e-4
            Sigma[singularity] = 4*r_s*rho_0/15
            Sigma[ordinary]    = rho_0*r_s*((2+s[ordinary]*s[ordinary])*X(s[ordinary])-3)/(1-s[ordinary]*s[ordinary])**2
            return Sigma
        else:
            if np.abs(s - 1.0) <= 1e-4:
                return 4*r_s*rho_0/15
            else:
                return rho_0*r_s*((2+s*s)*X(s)-3)/(1-s*s)**2
    
    def M_hernquist_2D(self,rho_0,r_s):
        s = self.r/r_s
        if np.abs(s - 1.0) <= 1e-4:
            return 2*np.pi*r_s**3*rho_0/3
        if s == np.inf:
            return 2*np.pi*r_s**3*rho_0
        else:
            s += 3e-4   # perturbation to avoid singularities in arcsech and arcsec
            return 2*np.pi*r_s**3*rho_0*s*s*(X(s)-1)/(1-s*s)
            
if __name__ == "__main__":
    
    sirius_dir = "/work/asl/AstrophysicalProject_Lensing"
    
    import matplotlib.pyplot as plt
    
    r = np.logspace(-2,1)
    R = np.logspace(-2,1)
    scale = 1.0
    M10 = 1.0
    rho10 = M10/(4.0/3.0*np.pi*1000)
    
    fig, ax =  plt.subplots(1,4,figsize = (48,10))
    ax[0].plot(r,sph_profile(r).power_law(M10/sph_profile(r[-1]).M_power_law(1.0,scale,2.0),scale,2.0)/rho10,label = "isothermal")
    ax[1].plot(r,sph_profile(r).two_power_law(M10/sph_profile(r[-1]).M_two_power_law(1.0,scale,1.0,3.0),scale,1.0,3.0)/rho10,label = "NFW")
    ax[2].plot(r,sph_profile(r).two_power_law(M10/sph_profile(r[-1]).M_two_power_law(1.0,scale,1.0,4.0),scale,1.0,4.0)/rho10,label = "Hernquist")
    ax[3].plot(r,sph_profile(r,dim="3D").sersic(M10/sph_profile(r[-1],dim="3D").M_sersic(1.0,scale,4.0),scale,4.0)/rho10,label = "Sersic")
    for i in range(4):
        ax[i].set_xscale("log")
        ax[i].set_yscale("log")
        ax[i].set_xlabel(r"$r/r_{\rm s}$",fontsize = 35)
    ax[0].set_ylabel(r"$\rho/\bar{\rho}$",fontsize = 35)
    fig.savefig(sirius_dir+"/Scripts/sph_dyn_lens/example_density.png",bbox_inches = "tight")
    
    fig, ax =  plt.subplots(1,4,figsize = (48,10))
    ax[0].plot(r,sph_profile(r,dim="2D").M_power_law(M10/sph_profile(r[-1],dim="2D").M_power_law(1.0,scale,2.0),scale,2.0),label = "isothermal")
    ax[1].plot(r,sph_profile(r,dim="2D").M_two_power_law(M10/sph_profile(r[-1],dim="2D").M_two_power_law(1.0,scale,1.0,3.0),scale,1.0,3.0),label = "NFW")
    ax[2].plot(r,sph_profile(r,dim="2D").M_two_power_law(M10/sph_profile(r[-1],dim="2D").M_two_power_law(1.0,scale,1.0,4.0),scale,1.0,4.0),label = "Hernquist")
    ax[3].plot(r,sph_profile(r,dim="2D").M_sersic(M10/sph_profile(r[-1],dim="2D").M_sersic(1.0,scale,4.0),scale,4.0),label = "Sersic")
    for i in range(4):
        ax[i].set_xscale("log")
        ax[i].set_yscale("log")
        ax[i].set_xlabel(r"$r/r_{\rm s}$",fontsize = 35)
    ax[0].set_ylabel(r"$M/M_{10}$",fontsize = 35)
    fig.savefig(sirius_dir+"/Scripts/sph_dyn_lens/example_mass.png",bbox_inches = "tight")